#include "precomp.h"

#ifdef _WIN32

#pragma hdrstop
#include <windows.h>
#include <psapi.h>
#include <tlhelp32.h>
#include <algorithm>
#include <winternl.h>

// 需要链接的库
#pragma comment(lib, "psapi.lib") // 用于GetModuleFileNameEx函数
#pragma comment(lib, "advapi32.lib")

struct ProcessInfo {
	DWORD pid; // 进程ID
	std::string exePath; // 可执行文件完整路径
	std::string cmdline;
};

using CProcessInfoList = std::vector<ProcessInfo>;

struct MemoryRegion {
	uintptr_t start_addr;
	uintptr_t end_addr;
	uintptr_t alloc_base;
	std::string permissions;
	std::string path;
	DWORD state;
	DWORD type;
	size_t size()
	{
		return end_addr - start_addr;
	}
};

using CMemoryRegionList = std::vector<MemoryRegion>;

struct Module {
	std::string path;
	uint64_t start, size;
};
using CModuleList = std::vector<Module>;

static CModuleList getProcessModules(HANDLE hProcess);
static std::string get_memory_path(const CModuleList& modules, LPCVOID a);
static std::string get_memory_permissions(DWORD protect);

// 提升进程权限（关键函数）
bool EnableDebugPrivilege()
{
	HANDLE hToken;
	TOKEN_PRIVILEGES tp;
	LUID luid;

	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
		return false;
	}

	if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &luid)) {
		CloseHandle(hToken);
		return false;
	}

	tp.PrivilegeCount = 1;
	tp.Privileges[0].Luid = luid;
	tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

	if (!AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(tp), NULL, NULL)) {
		CloseHandle(hToken);
		return false;
	}

	CloseHandle(hToken);
	return true;
}

#pragma comment(lib, "ntdll.lib")

// 将 wchar_t* (UTF-16) 转换为 std::string (UTF-8)
std::string WStringToUTF8(const wchar_t* wstr)
{
	if (!wstr || !*wstr)
		return "";

	// 计算需要的缓冲区大小
	int size_needed = WideCharToMultiByte(
		CP_UTF8, // 目标编码：UTF-8
		0, // 标志（一般填 0）
		wstr, // 输入的宽字符串
		-1, // 自动计算长度（直到 NULL 终止符）
		NULL, // 输出缓冲区（NULL 表示仅计算大小）
		0, // 输出缓冲区大小（0 表示仅计算）
		NULL, NULL // 默认字符和是否使用默认字符
	);

	if (size_needed <= 0) {
		return "";
	}

	// 分配缓冲区并转换
	std::string utf8_str(size_needed, 0);
	WideCharToMultiByte(
		CP_UTF8, 0, wstr, -1,
		&utf8_str[0], size_needed,
		NULL, NULL);

	return utf8_str;
}

// 获取进程命令行（返回 std::string）
static std::string GetProcessCommandLine(DWORD pid)
{
	HANDLE hProcess = OpenProcess(
		PROCESS_QUERY_INFORMATION | PROCESS_VM_READ,
		FALSE,
		pid);
	if (!hProcess) {
		return "";
	}

	// 获取进程基本信息
	PROCESS_BASIC_INFORMATION pbi;
	ULONG retLength;
	NTSTATUS status = NtQueryInformationProcess(
		hProcess,
		ProcessBasicInformation,
		&pbi,
		sizeof(pbi),
		&retLength);
	if (status != 0) {
		CloseHandle(hProcess);
		return "";
	}

	// 读取 PEB
	PEB peb;
	if (!ReadProcessMemory(hProcess, pbi.PebBaseAddress, &peb, sizeof(peb), NULL)) {
		CloseHandle(hProcess);
		return "";
	}

	// 读取 RTL_USER_PROCESS_PARAMETERS
	RTL_USER_PROCESS_PARAMETERS upp;
	if (!ReadProcessMemory(hProcess, peb.ProcessParameters, &upp, sizeof(upp), NULL)) {
		CloseHandle(hProcess);
		return "";
	}

	// 读取命令行（Unicode）
	WCHAR* cmdLine = (WCHAR*)malloc(upp.CommandLine.Length + 2);
	if (!ReadProcessMemory(hProcess, upp.CommandLine.Buffer, cmdLine, upp.CommandLine.Length, NULL)) {
		CloseHandle(hProcess);
		free(cmdLine);
		return "";
	}
	cmdLine[upp.CommandLine.Length / 2] = 0;
	CloseHandle(hProcess);
	std::string ret = WStringToUTF8(cmdLine);
	free(cmdLine);
	return ret;
}

CProcessInfoList EnumAllProcesses()
{
	std::vector<ProcessInfo> processes;
	PROCESSENTRY32 pe32 = { sizeof(PROCESSENTRY32) };

	// 创建系统进程快照
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hSnapshot == INVALID_HANDLE_VALUE) {
		return processes;
	}

	// 遍历进程列表
	if (Process32First(hSnapshot, &pe32)) {
		do {
			ProcessInfo info;
			info.pid = pe32.th32ProcessID;

			// 打开进程以获取完整路径
			HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pe32.th32ProcessID);
			if (hProcess != NULL) {
				char exePath[MAX_PATH] = { 0 };
				if (GetModuleFileNameEx(hProcess, NULL, exePath, MAX_PATH)) {
					info.exePath = exePath;
				}
				CloseHandle(hProcess);
			}
			info.cmdline = GetProcessCommandLine(info.pid);
			processes.push_back(info);
		} while (Process32Next(hSnapshot, &pe32));
	}

	CloseHandle(hSnapshot);
	return processes;
}

static CMemoryRegionList getProcessMemMap(HANDLE hProcess)
{
	// 获取系统内存范围
	SYSTEM_INFO systemInfo;
	GetSystemInfo(&systemInfo);
	LPVOID currentAddress = systemInfo.lpMinimumApplicationAddress;
	LPVOID maxAddress = systemInfo.lpMaximumApplicationAddress;
	CModuleList modules = getProcessModules(hProcess);
	CMemoryRegionList list;

	// 遍历内存区域
	while (currentAddress < maxAddress) {
		MEMORY_BASIC_INFORMATION memInfo;
		SIZE_T querySize = VirtualQueryEx(hProcess, currentAddress, &memInfo, sizeof(memInfo));
		if (querySize == 0) {
			// 查询失败，按页大小步进
			currentAddress = (LPVOID)((DWORD64)currentAddress + 0x1000);
			continue;
		}
		MemoryRegion r;
		r.start_addr = (uintptr_t)memInfo.BaseAddress;
		r.end_addr = (uintptr_t)memInfo.BaseAddress + memInfo.RegionSize;
		r.alloc_base = (uintptr_t)memInfo.AllocationBase;
		r.path = get_memory_path(modules, memInfo.BaseAddress);
		r.permissions = get_memory_permissions(memInfo.Protect);
		r.state = memInfo.State;
		r.type = memInfo.Type;
		list.push_back(r);
		// 移动到下一个内存区域
		currentAddress = (LPVOID)((DWORD64)memInfo.BaseAddress + memInfo.RegionSize);
	}
	return list;
}

// 内存搜索函数 - 在指定进程的地址空间中搜索特定字符串
CPointerList searchProcessString(pid_t pid, const char* searchString, size_t searchStringLength)
{
	std::vector<LPVOID> resultAddresses;
	if (searchStringLength <= 0)
		searchStringLength = strlen(searchString);

	// 打开目标进程
	HANDLE hProcess = OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, FALSE, pid);
	if (hProcess == NULL) {
		return resultAddresses; // 打开进程失败，返回空结果
	}

	// 获取系统信息以确定内存范围
	SYSTEM_INFO systemInfo;
	GetSystemInfo(&systemInfo);
	LPVOID currentAddress = systemInfo.lpMinimumApplicationAddress;
	LPVOID maxAddress = systemInfo.lpMaximumApplicationAddress;

	// 遍历进程的内存区域
	while (currentAddress < maxAddress) {
		MEMORY_BASIC_INFORMATION memInfo;
		if (VirtualQueryEx(hProcess, currentAddress, &memInfo, sizeof(memInfo)) == 0) {
			// 查询失败，移动到下一个地址
			currentAddress = (LPVOID)((DWORD64)currentAddress + 0x10000);
			continue;
		}

		// 检查内存区域是否可读取
		if ((memInfo.State == MEM_COMMIT) && ((memInfo.Protect & PAGE_READONLY) || (memInfo.Protect & PAGE_READWRITE) || (memInfo.Protect & PAGE_EXECUTE_READ) || (memInfo.Protect & PAGE_EXECUTE_READWRITE))) {

			// 分配缓冲区来存储读取的内存
			char* buffer = new char[memInfo.RegionSize];
			SIZE_T bytesRead;

			// 读取进程内存
			if (ReadProcessMemory(hProcess, memInfo.BaseAddress, buffer, memInfo.RegionSize, &bytesRead)) {
				// 在读取的内存中搜索字符串
				for (SIZE_T i = 0; i < bytesRead - searchStringLength; i++) {
					if (memcmp(buffer + i, searchString, searchStringLength) == 0) {
						// 找到匹配，计算在进程中的实际地址
						LPVOID matchAddress = (LPVOID)((DWORD64)memInfo.BaseAddress + i);
						resultAddresses.push_back(matchAddress);
					}
				}
			}

			delete[] buffer;
		}

		// 移动到下一个内存区域
		currentAddress = (LPVOID)((DWORD64)memInfo.BaseAddress + memInfo.RegionSize);
	}

	// 关闭进程句柄
	CloseHandle(hProcess);

	return resultAddresses;
}

static int lua_enumProcesses(lua_State* L)
{
	CProcessInfoList plist = EnumAllProcesses();
	const char* pattern = lua_tostring(L, 1);
	rtk::CRegexMatcher re;
	if (pattern) {
		if (re.compile(pattern, false) != 0) {
			luaL_error(L, "Pattern error @ '%s':%d.", pattern, re.erroffset);
		}
	}
	lua_createtable(L, plist.size(), 0);
	int i = 1;
	for (auto p : plist) {
		if (!pattern || re.match(p.exePath.c_str())) {
			lua_createtable(L, 0, 2);
			lua_pushliteral(L, "id");
			lua_pushinteger(L, p.pid);
			lua_rawset(L, -3);
			lua_pushliteral(L, "exe");
			lua_pushstring(L, p.exePath.c_str());
			lua_rawset(L, -3);
			lua_pushliteral(L, "cmdline");
			lua_pushstring(L, p.cmdline.c_str());
			lua_rawset(L, -3);
			lua_rawseti(L, -2, i);
			i++;
		}
	}
	return 1;
}

static int lua_openProcess(lua_State* L)
{
	int pid = lua_tointeger(L, 1);
	HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
	if (hProcess == NULL)
		return 0;
	lua_pushinteger(L, (lua_Integer)hProcess);
	return 1;
}

static int lua_closeHandle(lua_State* L)
{
	HANDLE h = (HANDLE)lua_tointeger(L, 1);
	if (h && h != INVALID_HANDLE_VALUE) {
		CloseHandle(h);
	}
	return 0;
}

static int lua_readProcessMemory(lua_State* L)
{
	HANDLE h = (HANDLE)lua_tointeger(L, 1);
	if (h == NULL)
		luaL_error(L, "Invalid process handle.");
	PVOID start = (PVOID)lua_tointeger(L, 2);
	size_t size = lua_tointeger(L, 3);
	size_t got = 0;
	void* buffer = malloc(size);
	if (!buffer)
		return 0;
	if (!ReadProcessMemory(h, start, buffer, size, &got)) {
		free(buffer);
		return 0;
	}
	lua_pushlstring(L, (char*)buffer, got);
	free(buffer);
	return 1;
}

static int lua_writeProcessMemory(lua_State* L)
{
	HANDLE h = (HANDLE)lua_tointeger(L, 1);
	if (h == NULL)
		luaL_error(L, "Invalid process handle.");
	PVOID start = (PVOID)lua_tointeger(L, 2);
	size_t size;
	const char* buffer = luaL_checklstring(L, 3, &size);
	size_t got;
	PVOID base = (PVOID)((unsigned long long)start & ~0xfff);
	DWORD oldProtect;
	VirtualProtectEx(h, base, 0x1000, PAGE_EXECUTE_READWRITE, &oldProtect);
	BOOL ok = WriteProcessMemory(h, start, buffer, size, &got);
	VirtualProtectEx(h, base, 0x1000, oldProtect, 0);
	if (!ok) {
		lua_pushnil(L);
		lua_pushinteger(L, GetLastError());
		return 2;
	}
	lua_pushinteger(L, got);
	return 1;
}

static std::string get_memory_permissions(DWORD protect)
{
	std::string perms = format_std_str("%08x", protect);
	auto check = [&protect, &perms](int bit, const char* sbit) {
		if (protect & bit) {
			protect &= ~bit;
			perms += sbit;
			perms += ' ';
		}
	};
	if (protect) {
		perms += "(";
		check(PAGE_EXECUTE, "x");
		check(PAGE_EXECUTE_READ, "rx");
		check(PAGE_EXECUTE_READWRITE, "rwx");
		check(PAGE_EXECUTE_WRITECOPY, "rxc");
		check(PAGE_NOACCESS, "na");
		check(PAGE_READONLY, "r");
		check(PAGE_READWRITE, "rw");
		check(PAGE_WRITECOPY, "rc");
		check(PAGE_GUARD, "guard");
		check(PAGE_NOCACHE, "ncache");
		check(PAGE_WRITECOMBINE, "wc");
		if (protect)
			perms += format_std_str("%x ", protect);
		if (perms[perms.size() - 1] == ' ')
			perms.resize(perms.size() - 1);
		perms += ")";
	}
	return perms;
}

static CModuleList getProcessModules(HANDLE hProcess)
{
	DWORD cbNeeded;
	std::vector<HMODULE> handles;
	CModuleList list;
	size_t N = 128;
	handles.resize(N);
	// 枚举进程所有模块（包括exe和dll）
	if (!EnumProcessModulesEx(hProcess, handles.data(), handles.size() * sizeof(HMODULE), &cbNeeded, LIST_MODULES_ALL)) {
		if (cbNeeded > N * sizeof(HMODULE)) {
			handles.reserve(cbNeeded / sizeof(HMODULE));
			if (!EnumProcessModulesEx(hProcess, handles.data(), handles.size() * sizeof(HMODULE), &cbNeeded, LIST_MODULES_ALL))
				return list;
			handles.resize(cbNeeded / sizeof(HMODULE));
		} else {
			return list;
		}
	} else {
		handles.resize(cbNeeded / sizeof(HMODULE));
	}

	for (size_t i = 0; i < handles.size(); i++) {
		HMODULE hModule = handles[i];
		MODULEINFO moduleInfo;
		// 获取模块的地址范围（BaseOfDll 和 SizeOfImage）
		if (!GetModuleInformation(hProcess, hModule, &moduleInfo, sizeof(moduleInfo))) {
			continue;
		}
		Module m;
		m.start = (uint64_t)moduleInfo.lpBaseOfDll;
		m.size = moduleInfo.SizeOfImage;
		char path[MAX_PATH] = { 0 };
		if (GetModuleFileNameExA(hProcess, hModule, path, MAX_PATH)) {
			m.path = path;
		}
		list.push_back(m);
	}

	std::sort(list.begin(), list.end(), [](const Module& a, const Module& b) -> bool { return a.start < b.start; });
	return list;
}

static int lua_getProcessModules(lua_State* L)
{
	pid_t pid = lua_tointeger(L, 1);
	lua_newtable(L);
	size_t idx = 1;
	// 打开目标进程（需要PROCESS_QUERY_INFORMATION和PROCESS_VM_READ权限）
	HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid);
	if (hProcess == NULL) {
		return 0;
	}
	CModuleList modules = getProcessModules(hProcess);
	for (auto m : modules) {
		lua_newtable(L);
		lua_pushliteral(L, "start");
		lua_pushinteger(L, m.start);
		lua_rawset(L, -3);
		lua_pushliteral(L, "size");
		lua_pushinteger(L, m.size);
		lua_rawset(L, -3);
		lua_pushliteral(L, "path");
		lua_pushlstring(L, m.path.c_str(), m.path.size());
		lua_rawset(L, -3);
		lua_rawseti(L, -2, idx);
		idx++;
	}
	return 1;
}

static std::string get_memory_path(const CModuleList& modules, LPCVOID a)
{
	size_t baseAddr = (size_t)a;
	for (auto m : modules) {
		if (m.start <= baseAddr && m.start + m.size > baseAddr)
			return m.path;
	}
	// 非模块内存区域（如堆、栈等）
	return "";
}

static int lua_getProcessMemMap(lua_State* L)
{
	pid_t pid = lua_tointeger(L, 1);
	lua_newtable(L);
	size_t idx = 1;

	// 打开目标进程（需要PROCESS_QUERY_INFORMATION和PROCESS_VM_READ权限）
	HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid);
	if (hProcess == NULL) {
		return 1; // 返回空表而非0，避免Lua栈不平衡
	}
	CMemoryRegionList list = getProcessMemMap(hProcess);
	for (MemoryRegion& r : list) {
		// 创建当前内存区域的信息表
		lua_newtable(L);

		// 起始地址（base address）
		lua_pushliteral(L, "start");
		lua_pushinteger(L, r.start_addr);
		lua_rawset(L, -3);

		// 结束地址（base + size）
		lua_pushliteral(L, "end");
		lua_pushinteger(L, r.end_addr);
		lua_rawset(L, -3);

		// 区域大小
		lua_pushliteral(L, "size");
		lua_pushinteger(L, r.size());
		lua_rawset(L, -3);

		// 内存对应的模块路径（如exe、dll）
		lua_pushliteral(L, "path");
		lua_pushlstring(L, r.path.c_str(), r.path.size());
		lua_rawset(L, -3);

		// 内存权限（如rwx、rw等）
		lua_pushliteral(L, "permission");
		lua_pushlstring(L, r.permissions.c_str(), r.permissions.size());
		lua_rawset(L, -3);

		lua_pushliteral(L, "alloc_base");
		lua_pushinteger(L, r.alloc_base);
		lua_rawset(L, -3);

		lua_pushliteral(L, "state");
		lua_pushstring(L, r.state == MEM_COMMIT ? "commit" : r.state == MEM_FREE ? "free"
				: r.state == MEM_RESERVE										 ? "reserved"
																				 : "");
		lua_rawset(L, -3);

		lua_pushliteral(L, "type");
		lua_pushstring(L, r.type == MEM_IMAGE ? "image" : r.type == MEM_MAPPED ? "mapped"
				: r.type == MEM_PRIVATE										   ? "private"
																			   : "");
		lua_rawset(L, -3);

		// 将当前区域信息插入结果表
		lua_rawseti(L, -2, idx);
		idx++;
	}

	CloseHandle(hProcess);
	return 1; // 返回结果表
}

/*
add memory operation funcs
module table is on top of stack on entry
*/
static int add_memory_ops(lua_State* L)
{
	EnableDebugPrivilege();
	regfunc(enumProcesses);
	regfunc(openProcess);
	regfunc(closeHandle);
	regfunc(readProcessMemory);
	regfunc(writeProcessMemory);
	regfunc(getProcessMemMap);
	regfunc(getProcessModules);
	return 0;
}

GAMEX_ADD_INIT_FUNC(add_memory_ops)

#ifdef GAMEX_TEST
int main(int c, char** v)
{
	test0(c, v);
	return 0;
}
#endif

#endif // _WIN32
